from langchain_community.document_loaders import (
    TextLoader,
    PythonLoader,
    JSONLoader,
    CSVLoader,
    UnstructuredMarkdownLoader
)
from langchain.schema import Document
import os
from typing import Dict, List, Optional


class CodeLoader:
    """代码文件加载器"""

    # 文件扩展名到加载器的映射
    EXTENSION_TO_LOADER = {
        ".py": PythonLoader,
        ".js": TextLoader,
        ".ts": TextLoader,
        ".java": TextLoader,
        ".json": JSONLoader,
        ".csv": CSVLoader,
        ".md": UnstructuredMarkdownLoader,
        ".txt": TextLoader,
    }

    @classmethod
    def get_loader_for_file(cls, file_path: str):
        """根据文件类型返回适当的加载器

        Args:
            file_path: 文件路径

        Returns:
            适用于该文件类型的加载器实例
        """
        _, ext = os.path.splitext(file_path)

        if ext not in cls.EXTENSION_TO_LOADER:
            # 默认使用TextLoader
            return TextLoader(file_path)

        loader_class = cls.EXTENSION_TO_LOADER[ext]

        # 特殊处理JSON加载器，它需要jq参数
        if loader_class == JSONLoader:
            return JSONLoader(file_path, jq_schema='.', text_content=False)

        return loader_class(file_path)

    @classmethod
    def load_project(cls, project_root: str, exclude_dirs: Optional[List[str]] = None) -> Dict[str, List[Document]]:
        """加载整个项目的所有文件

        Args:
            project_root: 项目根目录
            exclude_dirs: 要排除的目录列表

        Returns:
            文件路径到文档列表的映射
        """
        if exclude_dirs is None:
            exclude_dirs = ['.git', '__pycache__', 'venv', 'env', '.venv', 'node_modules']

        result = {}

        for root, dirs, files in os.walk(project_root):
            # 排除指定目录
            dirs[:] = [d for d in dirs if d not in exclude_dirs]

            for file in files:
                file_path = os.path.join(root, file)
                _, ext = os.path.splitext(file)

                # 只处理我们支持的文件类型
                if ext in cls.EXTENSION_TO_LOADER or ext == "":
                    try:
                        loader = cls.get_loader_for_file(file_path)
                        documents = loader.load()
                        # 使用相对路径作为键
                        rel_path = os.path.relpath(file_path, project_root)
                        result[rel_path] = documents
                    except Exception as e:
                        print(f"加载文件 {file_path} 时出错: {e}")

        return result

    @classmethod
    def load_file(cls, file_path: str) -> List[Document]:
        """加载单个文件

        Args:
            file_path: 文件路径

        Returns:
            文档对象列表
        """
        try:
            loader = cls.get_loader_for_file(file_path)
            return loader.load()
        except Exception as e:
            print(f"加载文件 {file_path} 时出错: {e}")
            return []
